Benchmark Framework Design

This document describes the overall design of the Liger-Kernel benchmark suite, including its two benchmark dimensions, the shared infrastructure, and the phased implementation plan.

1. Benchmark Dimensions

Every operator should ideally be benchmarked along two orthogonal dimensions:

D1: Non-model dimension sweep (implemented)

Sweep non-model dimensions (e.g. sequence length, BT) with a fixed model config selected via --model. This is the default behavior for all benchmark scripts.

x_values:  [1024, 2048, 4096, 8192, ...]     (seq_len or BT)
fixed:     model=llama_3_8b via --model       (hidden_size=4096, intermediate_size=14336, ...)
output:    line chart — speed/memory vs token length

D2: Model dimension sweep (implemented)

Sweep model architecture dimensions (e.g. hidden_size, or discrete model configs from MODEL_REGISTRY) with a fixed token count set via --bt. This reveals how kernel performance compares across different model architectures at the same input scale.

x_values:  [llama_2_7b, llama_3_8b, ...]     (discrete model configs)
fixed:     BxT via --bt                       (determined to be safe across all configs)
output:    speedup or throughput bar chart per model config

2. D2 Design Choices

Following the maintainer discussion, we evaluated three approaches:

Decision: C as the primary approach, with A as optional enrichment for ops where single-parameter scaling is important.

Rationale:

• C uses real-world model architectures, making results directly meaningful.
• C naturally extends the existing MODEL_REGISTRY infrastructure.
• C produces clean bar charts (speedup/throughput) that align across ops.
• A can be layered on later for specific ops where parameter-level scaling trends matter.

3. Universal Token Length for D2

For D2 benchmarks, we need a fixed token-length that is safe (no OOM) across all model configs and all operators.

Strategy

1. Coupled tokens: define a (batch_size, seq_len) pair per model config, e.g. (B=2, T=1024) for llama_2_7b, (B=1, T=1024) for llama_3_8b. This allows adapting to each model's memory footprint while keeping the comparison fair (same total token count or same seq_len).
2. Safety via probe: before running D2 benchmarks, run estimate_kernel_peak_memory for each (model_config, token_config) pair. If any config would OOM, reduce token count automatically.
3. Forward compatibility: when new ops are added, the probe mechanism ensures safety without manual tuning. If a new op cannot fit any reasonable token count for a given model, the framework skips that data point with a warning rather than crashing.

Proposed CLI

# D1 (existing): token-length sweep with fixed model
python benchmark_geglu.py --model llama_3_8b

# D2 (new): model-config sweep with fixed token length
python benchmark_geglu.py --sweep-mode model_config --bt 2048

The --sweep-mode flag selects the dimension. Default remains token_length (D1) for backward compatibility.

4. Infrastructure Changes

4.1 New config type

@dataclass(frozen=True)
class ModelConfigSweepConfig:
    """Config for D2 benchmarks that sweep across model configs."""
    model_configs: List[ModelConfig]       # models to benchmark
    bt: int                                 # fixed batch * seq_len
    batch_size: int                         # safe batch size
    seq_len: int                            # safe seq_len

4.2 New helper

def compute_model_config_sweep_config(
    model_configs: List[ModelConfig],
    probe_fn_factory: Callable[[ModelConfig, int], Callable[[], torch.Tensor]],
    bt: int = 2048,
    memory_utilization: float = 0.4,
) -> ModelConfigSweepConfig:
    """Find safe (batch_size, seq_len) that works across all model configs.

    For each model config, runs probe_fn_factory(model_config, bt) to measure
    peak memory, then picks the most conservative batch_size / seq_len.
    """
    ...

4.3 Script-level changes

Each benchmark script gains a model-config sweep code path gated by --sweep-mode:

if args.sweep_mode == "model_config":
    configs = [MODEL_REGISTRY[name] for name in MODEL_REGISTRY]
    sweep = compute_model_config_sweep_config(configs, probe_fn_factory=..., bt=args.bt)
    # x_values = model config indices
    # extra_benchmark_configs = contains all model configs
    ...
else:
    # existing token-length sweep logic
    ...

4.4 Visualization

D2 results produce grouped bar charts (speedup or throughput) rather than line charts:

• x-axis: model config names (e.g. llama_2_7b, llama_3_8b)
• bars: kernel providers (liger vs huggingface/torch)
• y-axis: speedup ratio or throughput (tokens/s)

5. Phased Implementation Plan

Phase 1: Foundation (current PR)

Status: complete

• ModelConfig and MODEL_REGISTRY with canonical model profiles
• estimate_kernel_peak_memory() for runtime memory probing
• compute_seq_len_sweep_config() for D1 seq_len sweeps (non-model dimension)
• compute_hidden_size_sweep_config() for D2 hidden_size sweeps (model dimension)
• run_speed_benchmark() / run_memory_benchmark() shared helpers
• --model CLI argument
• Refactored benchmark_geglu.py, benchmark_swiglu.py, benchmark_dyt.py
• BENCHMARK_GUIDELINES.md contributor guide

Phase 2: Model-config sweep (D2)

Status: complete

• Add ModelConfigSweepConfig dataclass
• Implement compute_model_config_sweep_config() with cross-model probe
• Add --sweep-mode and --bt CLI arguments to parse_benchmark_script_args()
• Add model-config sweep code path to benchmark_geglu.py as reference implementation
• Model-config sweep code path ported to benchmark_swiglu.py and benchmark_dyt.py
• Validate on at least 2 devices (CUDA + NPU) to confirm OOM safety
• Update BENCHMARK_GUIDELINES.md with D2 instructions

Phase 3: Rollout and visualization

Status: in progress

• Add bar chart / speedup visualization for D2 results
• Integrate into CI workflows

Phase 3 Kernel Rollout Tracking

Already refactored (D1 + D2):

• benchmark_geglu.py
• benchmark_swiglu.py
• benchmark_dyt.py

Norm-like kernels (input: BT × hidden_size):

• benchmark_rms_norm.py
• benchmark_layer_norm.py
• benchmark_group_norm.py
• benchmark_fused_add_rms_norm.py
• benchmark_poly_norm.py

Loss kernels (input involves vocab_size or similar):

• benchmark_cross_entropy.py
• benchmark_fused_linear_cross_entropy.py
• benchmark_jsd.py
• benchmark_fused_linear_jsd.py
• benchmark_kl_div.py
• benchmark_tvd.py
• benchmark_distill_cosine_loss.py
• benchmark_distill_jsd_loss.py

RLHF/alignment loss kernels:

• benchmark_dpo_loss.py
• benchmark_cpo_loss.py
• benchmark_orpo_loss.py
• benchmark_simpo_loss.py
• benchmark_kto_loss.py
• benchmark_grpo_loss.py

Positional encoding kernels:

• benchmark_rope.py
• benchmark_qwen2vl_mrope.py
• benchmark_llama4_rope.py

Activation / misc kernels:

• benchmark_softmax.py
• benchmark_sparsemax.py
• benchmark_embedding.py
• benchmark_relu_squared.py

Attention kernels:

• benchmark_fused_neighborhood_attention.py
• benchmark_sparse_multi_token_attention.py
• benchmark_multi_token_attention.py
• benchmark_attn_res.py

Other:

• benchmark_tiled_mlp.py
• benchmark_mhc.py
• benchmark_mhc_lm.py

6. Directory Structure

benchmark/
├── data/
│   └── all_benchmark_data.csv
├── scripts/
│   ├── benchmark_model_configs.py      # ModelConfig, MODEL_REGISTRY, helpers
│   ├── utils.py                        # run_benchmarks, CSV, CLI
│   ├── benchmark_geglu.py              # D1 + D2
│   ├── benchmark_swiglu.py             # D1 + D2
│   ├── benchmark_dyt.py                # D1 + D2
│   └── ...
├── visualize/                          # (Phase 3) chart generation
│   └── ...
└── BENCHMARK_GUIDELINES.md             # contributor guide